import numpy as np
import torch
from torch import nn


# 非常耗时，CPU操作太多导致GPU等待
# TODO: 如何使用cuda编程来优化
class DistanceLossV1(nn.Module):
    def __init__(self):
        super(DistanceLossV1, self).__init__()

    def forward(self, points, lengths, feature_size):
        b, _, h, w = feature_size
        max_len = max(lengths)
        assert points.size(0) == b and lengths.size(0) == b and points.size(1) == max_len
        loss = torch.tensor(0.0, dtype=points.dtype).to(points.device)
        for i in range(b):
            char_num = torch.tensor(lengths[i] - 1, dtype=points.dtype).to(points.device)  # lengths[i]包含结束符EOS
            char_width = w / char_num  # 将特征等分成char_num份
            for j in range(0, max_len):
                if j < char_num:  # 对于前char_num个点，需要对应到每个字符上
                    # 左右边界的计算还可以再灵活点
                    left_margin = char_width * j
                    right_margin = char_width * (j + 1)
                    if left_margin <= points[i][j][0] <= right_margin:  # points[i][j][0] 为x坐标
                        loss += 0
                    elif points[i][j][0] < left_margin:
                        loss += left_margin - points[i][j][0]
                    else:
                        loss += points[i][j][0] - right_margin
                else:  # char_num后面的点都是多余的，应该放到特征右边缘的地方
                    left_margin = w - 1 - char_width
                    right_margin = w - 1
                    if points[i][j][0] < left_margin:
                        loss += left_margin - points[i][j][0]
        loss = loss / b
        return loss



class DistanceLossV2(nn.Module):
    def __init__(self):
        super(DistanceLossV2, self).__init__()
        self.l1_loss = nn.L1Loss()
        

    def forward(self, points, lengths, feature_size):
        b, _, h, w = feature_size
        max_len = max(lengths)
        assert points.size(0) == b and lengths.size(0) == b and points.size(1) == max_len
        loss = torch.tensor(0.0, dtype=points.dtype).to(points.device)
        # target = torch.zeros((b, max_len), dtype=points.dtype).to(points.device)
        target = []
        for i in range(b):
            char_num = lengths[i].detach().item() - 1  # lengths[i]包含结束符EOS
            char_width = w / char_num  # 将特征等分成char_num份
            head_x = np.arange(char_width / 2, w - 1, char_width)
            # head_y = np.clip(np.random.normal((h - 1) / 2, 0.5, len(head_x)), (h - 1) / 2 - 1, (h - 1) / 2 + 1)  # 随机生成复合正态分布的序列作为y坐标（均值为(h-1)/2，方差为0.5，限定范围在[0.5, 2.5]）
            padding_x = np.ones(int((max_len - len(head_x)).detach().item())) * head_x[-1]
            # padding_y = np.ones(int((max_len - len(head_x)).detach().item())) * ((h - 1) / 2)
            # target.append([np.append(head_x, padding_x), np.append(head_y, padding_y)])
            target.append([np.append(head_x, padding_x), np.ones(max_len.detach().item()) * ((h - 1) / 2)])
        target = torch.from_numpy(np.array(target)).float().to(points.device)
        target = target.permute(0, 2, 1)
        loss = self.l1_loss(points, target)
        return loss


if '__name__' == '__main__':
    import numpy as np
    np.random.standard_normal()
